package com.augrain.microframework.trace.instrument.rabbitmq;

import com.augrain.microframework.trace.TraceConstant;
import com.augrain.microframework.trace.TraceProperty;
import com.augrain.microframework.trace.common.TraceLabel;
import com.augrain.microframework.trace.context.TraceContext;
import com.augrain.microframework.trace.util.MDCUtils;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessagePostProcessor;
import org.springframework.amqp.core.MessageProperties;

/**
 * 解析header中的traceId，spanId
 *
 * @author biaoy
 * @since 2023/08/10
 */
public class AfterMessagePostProcessor implements MessagePostProcessor {

    private final TraceProperty traceProperty;

    public AfterMessagePostProcessor(TraceProperty traceProperty) {
        this.traceProperty = traceProperty;
    }

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {
        MessageProperties messageProperties = message.getMessageProperties();

        TraceLabel traceLabel;
        if (messageProperties.getHeader(TraceConstant.TRACE_ID) != null) {
            String traceId = messageProperties.getHeader(TraceConstant.TRACE_ID);
            String spanId = messageProperties.getHeader(TraceConstant.SPAN_ID);
            String preApp = messageProperties.getHeader(TraceConstant.FROM_APP);
            String preIp = messageProperties.getHeader(TraceConstant.FROM_IP);
            traceLabel = new TraceLabel(preApp, preIp, traceId, spanId);
        } else {
            // 没有透传
            traceLabel = new TraceLabel(null, null, null, null);
        }
        // 设置到线程上下文中
        TraceContext.setTraceLabel(traceLabel);
        // 设置到MDC中
        MDCUtils.put(traceLabel, traceProperty);

        return message;
    }
}
